#include <stdio.h>
#include <stdlib.h>
#include "prime.h"

struct Pair {
    int p1, p2;
};

/* n = 2^s * d. This method returns (s, d). */
struct Pair factor_twos(int n) {
    int ret = 0;
    struct Pair p;
    p.p1 = 0;
    p.p2 = n;
    while (p.p2 % 2 == 0) {
        p.p1 += 1;
        p.p2 = p.p2 >> 1;
    }

    return p;
}

/* Bruce Shchneier's Right-to-left binary method for computing
 * b^e (mod m)
 *
 * Easy on memory, doesn't use Math.pow, and only requires log(e)
 * multiplications. The multiplications are also usually kept small by
 * the modulus.
 */
int pow_mod(int b, int e, int m) {
    int ret = 1;

    while (e > 0) {
        if (e % 2 == 1)
            ret = (ret * b) % m;
        e = e >> 1;
        b = (b * b) % m;
    }

    return ret;
}

/* Miller--Rabin deterministic primality test
 *
 * fairly fast prime test thanks to an assumption we can make
 */
int is_prime(int n, int k) {
    int s, d, s_temp;
    int temp_ad, temp_a2rd;
    int i, go;
    struct Pair p;
    int lim;

    p = factor_twos(n - 1);
    s = p.p1;
    d = p.p2;


    /* Pomerance, Selfridge, Wagstaff, and Jaeschke proved that for
     * n < 1,373,653 it is enough to test i = 2 and 3
     *
     * Because we do not expect to use 500,000+ particles, we can cut this
     * test short
     */ 
    for (i = 2; i < 3; i++) {
        go = 0;
        for (s_temp = 0; s_temp < s; s_temp++) {
            temp_ad = pow_mod(i, d, n);
            temp_a2rd = pow_mod(i, d * (1 << s_temp), n);
            /* if i^d =\= 1 (mod n) AND i^((2^s_temp) * d) =\= -1 (mod n)
             * for all s_temp, then n is composite
             *
             * or we take the contradiction:
             * if any s_temp exists such that
             *
             * i^d == 1 (mod n) OR i^((2^s_temp) * d) == -1 (mod n)
             *
             * then we can shortcircuit the rest of s_temp
             */
            if (temp_ad == 1 || temp_ad == -n + 1 ||
                temp_a2rd == -1 || temp_a2rd == n - 1) {
                go = 1;
                break;
            }
        }

        if (!go)
            return 0;
    }

    return 1;
}

int next_prime(int n) {
    int ret;

    /* Bertrand's Theorem states that for all n > 3, 
     * there exists a prime p such that n < p < n * 2 - 2
     */
    ret = n + 1;
    if (ret % 2 == 0)
        ++ret;

    while (ret < n * 2 - 2) {
        if (is_prime(ret, 2)) {
            return ret;
        }

        ret += 2;
    }

    return ret;
}
